/*******************************************************************************
 * Copyright 2010 DemandSpot
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 ******************************************************************************/
package com.demandspot.aws;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.util.ArrayList;
import java.util.Calendar;
import java.util.Date;
import java.util.Deque;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.zip.GZIPInputStream;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;

import org.apache.log4j.Logger;
import org.apache.log4j.lf5.util.StreamUtils;
import org.jets3t.apps.synchronize.Synchronize;
import org.jets3t.service.Constants;
import org.jets3t.service.Jets3tProperties;
import org.jets3t.service.S3Service;
import org.jets3t.service.S3ServiceException;
import org.jets3t.service.acl.AccessControlList;
import org.jets3t.service.acl.GroupGrantee;
import org.jets3t.service.acl.Permission;
import org.jets3t.service.impl.rest.httpclient.RestS3Service;
import org.jets3t.service.io.BytesProgressWatcher;
import org.jets3t.service.model.S3Bucket;
import org.jets3t.service.model.S3Object;
import org.jets3t.service.multithread.S3ServiceSimpleMulti;
import org.jets3t.service.security.AWSCredentials;


public class AWSS3Service {
	private S3Service s3Service;
	private S3ServiceSimpleMulti simpleMulti;
	
	private static final Logger logger = Logger.getLogger(AWSS3Service.class);
	private static final int BATCH_SIZE = 100;
	
	public AWSS3Service() throws FileNotFoundException, IOException, AWSS3ServiceException {		
		this(AWSS3Service.class.getResourceAsStream("/conf/aws.properties"));
	}

	public AWSS3Service(InputStream propInputStream) throws IOException, AWSS3ServiceException {
		Properties props = new Properties();
		props.load(propInputStream);
		init(props);
	}
	
	public AWSS3Service(File propsFile) throws FileNotFoundException, IOException, AWSS3ServiceException {
		Properties props = new Properties();
		props.load(new FileInputStream(propsFile));
		init(props);
	}
	
	public AWSS3Service(URL propsURL) throws FileNotFoundException, IOException, AWSS3ServiceException {
		this(new File(propsURL.getFile()));
	}
	
	public AWSS3Service(Properties props) throws AWSS3ServiceException {
		init(props);
	}

	public AWSS3Service(String accessKey, String secretKey) throws AWSS3ServiceException {
		Properties props = new Properties();
		props.put(AWSConsts.AWS_ACCESS_KEY, accessKey);
		props.put(AWSConsts.AWS_SECRET_KEY, secretKey);
		init(props);
	}
	
	private void init(Properties props) throws AWSS3ServiceException {
		String accessKey = props.getProperty(AWSConsts.AWS_ACCESS_KEY);
		String secretKey = props.getProperty(AWSConsts.AWS_SECRET_KEY);
		
		accessKey = System.getProperty(AWSConsts.AWS_ACCESS_KEY, accessKey);
		secretKey = System.getProperty(AWSConsts.AWS_SECRET_KEY, secretKey);
		
		AWSCredentials awsCredentials = new AWSCredentials(accessKey, secretKey);
		
		try {
			s3Service = new RestS3Service(awsCredentials);	
			simpleMulti = new S3ServiceSimpleMulti(s3Service);
		} catch (S3ServiceException e) {
			throw new AWSS3ServiceException(e);
		}
	}

	/**
	 * Load a file to S3 to a given bucket. In-case that the bucket doesn't exist the method will try to create a new one.
	 * if the given file is a directory the method will upload all the directory content
	 * 
	 * @param bucketName
	 * @param file
	 * @param key
	 * @throws AWSS3ServiceException 
	 */
	public void putFile(String bucketName, String key, File file, boolean publicReadAllowed) throws AWSS3ServiceException {
		key = key == null? "" : key.trim();
		try {
			S3Bucket bucket = s3Service.getBucket(bucketName);
			
			//bucket doesn't exist - create a new bucket
			if (bucket == null) {
				bucket = s3Service.createBucket(bucketName);
			}

			AccessControlList bucketAcl = s3Service.getBucketAcl(bucket);
			if (publicReadAllowed) {				
				bucketAcl.grantPermission(GroupGrantee.ALL_USERS, Permission.PERMISSION_READ);
				bucket.setAcl(bucketAcl);
				s3Service.putBucketAcl(bucket);
			}

			
			//special handling for directories
			if (file.isDirectory()) {
				List<S3Object> objects = new ArrayList<S3Object>();
				Deque<File> dirStack = new LinkedList<File>();
				
				dirStack.push(file);
				int fullPrefixSize = file.getAbsolutePath().length();
				
				while(!dirStack.isEmpty()) {
					File dir = dirStack.pop();
					
					String prefix = dir.getAbsolutePath().substring(fullPrefixSize).replace('\\', '/');
					
					prefix = (key.isEmpty()? "" : key) + prefix;
					prefix += (prefix.isEmpty()? "" : '/');
					
					if (!prefix.isEmpty() && prefix.charAt(0) == '/') {
						prefix = prefix.substring(1);
					}
					
					File[] files = dir.listFiles();
										
					for (File file2 : files) {
						if (file2.isDirectory()) {
							dirStack.push(file2);
						} else {
							String objectKey = prefix + file2.getName();
							S3Object object = new S3Object(bucket, file2);
							object.setKey(objectKey);
							object.setAcl(bucketAcl);
							object.setLastModifiedDate(new Date(file.lastModified()));
							objects.add(object);
//							logger.info("Uploading " + file2);
//							s3Service.putObject(bucket, object);
						}
					}
				}
				
				logger.info("Uploading " + objects.size() + " files...");	
				for (int i = 0; i < objects.size() / BATCH_SIZE; i++) {
					int start = i*BATCH_SIZE;
					int end = (i+1)*BATCH_SIZE;
					if (end >= objects.size()) {
						end = objects.size()-1;
					}
					List<S3Object> partial = objects.subList(start, end);
					simpleMulti.putObjects(bucket, partial.toArray(new S3Object[0]));
					logger.info("Uploaded " + end + " files");
				}
				
			} else {
				S3Object object = new S3Object(bucket, file);
				String objectKey = (key.isEmpty()? "" : (key + '/')) + file.getName();
				object.setKey(objectKey);
				object.setAcl(bucketAcl);
				object.setLastModifiedDate(new Date(file.lastModified()));
				logger.info("Uploading " + file);
				s3Service.putObject(bucket, object);
			}
		} catch (Exception e) {
			throw new AWSS3ServiceException(e);
		}
	}
	
	/**
	 * Lists the objects in a bucket matching a prefix or all objects if prefix is null.
	 * 
	 * @param bucketName
	 * @param prefix
	 * @param delimiter
	 * @return list of files
	 * @throws AWSS3ServiceException
	 */
	public List<String> listFiles(String bucketName, String prefix, String delimiter) throws AWSS3ServiceException {
		S3Bucket bucket = null;
		try {
			bucket = s3Service.getBucket(bucketName);
			if (bucket != null) {
				S3Object[] objects = new S3Object[0];
				
				if (prefix == null) {
					objects = s3Service.listObjects(bucket);
				} else {
					objects = s3Service.listObjects(bucket, prefix, delimiter);
				}
				
				List<String> filesList = new ArrayList<String>();
				
				for (S3Object object : objects) {
					filesList.add(object.getKey());
				}
				
				return filesList;
				
			} else {
				throw new AWSS3ServiceException("Unknown bucket: " + bucketName);
			}
		} catch (Exception e) {
			throw new AWSS3ServiceException(e);
		}
	}
	
	/**
	 * 
	 * @param bucketName
	 * @param key
	 * @return
	 * @throws AWSS3ServiceException
	 */
	public InputStream getFile(String bucketName, String key) throws AWSS3ServiceException {
		S3Bucket bucket = null;
		try {
			bucket = s3Service.getBucket(bucketName);
			
			if (bucket != null) {
				S3Object object = s3Service.getObject(bucket, key);
				
				if (object != null) {
					InputStream toreturn = object.getDataInputStream();
					if (key.endsWith(".gz")) {
						toreturn = new GZIPInputStream(toreturn);
					}
					if (key.endsWith(".zip")) {
						toreturn = new ZipInputStream(toreturn);
					}
					return toreturn;
				} else {
					throw new AWSS3ServiceException("Unknown key: " + key);
				}
			} else {
				throw new AWSS3ServiceException("Unknown bucket: " + bucketName);
			}
		} catch (Exception e) {
			throw new AWSS3ServiceException(e);
		}
	}
	
	public InputStream getFile(String path) throws AWSS3ServiceException {
		String[] bucketAndKey = path.split(":");
		if (bucketAndKey.length != 2) {
			throw new IllegalArgumentException("path must be in the form bucket:key");
		}
		return getFile(bucketAndKey[0], bucketAndKey[1]);
	}

	public File getTempFile(String path) throws AWSS3ServiceException, IOException {
		File file = File.createTempFile("s3Tmp_", System.currentTimeMillis() + "");
		InputStream input = getFile(path);
		if (input instanceof ZipInputStream) {
			//create a directory for storing the files
			file.delete();
			file.mkdirs();
			ZipInputStream zinput = (ZipInputStream) input;
			ZipEntry entry = null;
			while ((entry = zinput.getNextEntry()) != null) {
				File entryFile = new File(file, entry.getName());
				if (entry.isDirectory()) {
					entryFile.mkdir();
				}
				else {
					FileOutputStream fout = new FileOutputStream(entryFile);
					for (int c = zinput.read(); c != -1; c = zinput.read()) {
						fout.write(c);
					}
					fout.close();
				}
		        zinput.closeEntry();				
			}			
		}
		else {
			FileOutputStream output = new FileOutputStream(file);
			StreamUtils.copy(input, output);
			output.close();						
		}
		input.close();
		return file;
	}
	
	private S3Bucket validateBucket(String bucketName) throws AWSS3ServiceException {
		try {
			S3Bucket bucket = s3Service.getBucket(bucketName);
			if (bucket == null) {
				throw new AWSS3ServiceException("Unknown bucket: " + bucketName);
			}
			return bucket;
		} catch (Exception e) {
			throw new AWSS3ServiceException(e);
		}
	}
	
	private S3Object validateObjectKey(S3Bucket bucket, String objectKey) throws AWSS3ServiceException {
		try {
			S3Object object = s3Service.getObject(bucket, objectKey);
			if (object == null) {
				throw new AWSS3ServiceException("Unknown object: " + objectKey);
			}
			return object;
		} catch (Exception e) {
			throw new AWSS3ServiceException(e);
		}
		
	}
	
	public void deleteFile(String bucketName, String key) throws AWSS3ServiceException {
		try {
			s3Service.deleteObject(bucketName, key);
		} catch (S3ServiceException e) {
			throw new AWSS3ServiceException(e);
		}
	}
	
	/**
	 * Move objects between 2 buckets
	 * @param key orig object key
	 * @param destKey dest object key
	 * @param srcBucket src bucket
	 * @param destBucket dest bucket
	 * @throws AWSS3ServiceException
	 */
	public void moveFile(String key, String destKey, String srcBucket, String destBucket) throws AWSS3ServiceException{
		try {
			S3Bucket bucket = validateBucket(srcBucket);
			S3Bucket dstBucket = validateBucket(destBucket);
			validateObjectKey(bucket, key);
			S3Object newObject = new S3Object(dstBucket, destKey);
			
			s3Service.moveObject(srcBucket, key, destBucket, newObject, false);
			
		} catch (AWSS3ServiceException e) {
			throw e;
		} catch (S3ServiceException e) {
			throw new AWSS3ServiceException(e);
		}
	}
	
	/**
	 * Synchronize a local folder with S3. Delete files that are not present locally, add new files and overwrite modified files.
	 * @param dir The local directory to sync
	 * @param bucketName The bucket name to write to (must exists beforehand)
	 * @param publicReadAllowed if <code>true</code> public read access is allowed
	 * @return
	 * @throws Exception
	 */
	public BytesProgressWatcher uploadFolder(File dir, String bucketName, boolean publicReadAllowed) throws Exception {
		String aclString = publicReadAllowed ? "PUBLIC_READ" : "PRIVATE";
		
		Jets3tProperties props = Jets3tProperties.getInstance(Constants.JETS3T_PROPERTIES_FILENAME);
		
		final Synchronize sync = new Synchronize(s3Service, 
				true, /*doAction*/ 
				false, /*isQuiet*/
				false, /*isNoProgress*/
				false, /*isForce*/
				false, /*isKeepFiles*/
				false, /*isNoDelete*/
				false, /*isMoveEnabled*/
				true, /*isBatchMode*/
				false, /*isSkipMetadata*/
				false, /*isGzipEnabled*/
				false, /*isEncryptionEnabled*/
				3 /*REPORT_LEVEL_ALL*/, 
				props);
		
		if (dir.isDirectory()) {
			Map<String, File> filesMap = new HashMap<String, File>();
			long totalBytes = 0;
			buildMap(filesMap, dir, "", totalBytes);
			S3Bucket bucket = validateBucket(bucketName);		
			BytesProgressWatcher progressWatcher = new BytesProgressWatcher(totalBytes);
			sync.uploadLocalDirectoryToS3(filesMap, bucket, dir.getName(), aclString, progressWatcher);			
			return progressWatcher;
		}
		else {		
			throw new IllegalArgumentException("File must be a directory");
		}
	}

	public void downloadFolder(File dir, String bucketName, String prefix) throws Exception {
		if (!dir.isDirectory()) {
			throw new IllegalArgumentException("File must be a directory");
		}
		String[] parts = prefix.split("/");
		File fl = dir;
		for (String part : parts) {
			 fl = new File(fl, part);			 
		}
		fl.mkdirs();
		S3Bucket bucket = null;
		try {
			bucket = s3Service.getBucket(bucketName);
			if (bucket != null) {				
				S3Object[] objects = s3Service.listObjects(bucket, prefix, "/");
				
				for (S3Object object : objects) {
					object = s3Service.getObject(bucket, object.getKey());
					InputStream stream = object.getDataInputStream();
					File file = new File(dir, object.getKey());					
					FileOutputStream fstream = new FileOutputStream(file);
					for (int c = stream.read(); c != -1; c = stream.read()) {
						fstream.write(c);
					}
					stream.close();
					fstream.close();						
				}
			} else {
				throw new AWSS3ServiceException("Unknown bucket: " + bucketName);
			}
		} catch (Exception e) {
			throw new AWSS3ServiceException(e);
		}		
	}
		
	private void buildMap(Map<String, File> filesMap, File dir, String key, long totalBytes) {		
		File[] files = dir.listFiles();
		for (File file : files) {
			totalBytes += file.length();
			String newKey = (key.length() > 0 ? key + "/" :  "") + file.getName();
			if (file.isDirectory()) {
				buildMap(filesMap, file, newKey, totalBytes);
			}
			else {
				filesMap.put(newKey, file);
			}
		}		
	}

	/**
	 * Create a read URL for a file for 10 years
	 * @param bucketName the bucket where the file is placed
	 * @param key the file object key
	 * @throws S3ServiceException
	 */
	public String createSignedGetUrl(String bucketName, String key) throws AWSS3ServiceException {
		Calendar cal = Calendar.getInstance();
		cal.add(Calendar.YEAR, 10);
		try {
			return S3Service.createSignedGetUrl(bucketName, key, s3Service.getAWSCredentials(), cal.getTime());
		} catch (S3ServiceException e) {
			throw new AWSS3ServiceException(e);
		}		
	}
	
}
